/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package edu.yonsei.iwb.smarttv.ner;

import java.io.IOException;
import java.util.Set;
import java.util.Arrays;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.ComplexExplanation;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Searcher;
import org.apache.lucene.search.Similarity;
import org.apache.lucene.search.function.CustomScoreProvider;
import org.apache.lucene.search.function.ValueSourceQuery;
import org.apache.lucene.util.ToStringUtils;

/**
 * Query that sets document score as a programmatic function of several (sub) scores:
 * <ol>
 *    <li>the score of its subQuery (any query)</li>
 *    <li>(optional) the score of its ValueSourceQuery (or queries).
 *        For most simple/convenient use cases this query is likely to be a 
 *        {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQuery}</li>
 * </ol>
 * Subclasses can modify the computation by overriding {@link #getCustomScoreProvider}.
 * 
 * @lucene.experimental
 */
public class CustomOrderScoreQuery extends Query {

    private Query subQuery;
    private ValueSourceQuery[] valSrcQueries; // never null (empty array if there are no valSrcQueries).
    private boolean strict = false; // if true, valueSource part of query does not take part in weights normalization.  

    /**
     * Create a CustomScoreQuery over input subQuery.
     * @param subQuery the sub query whose scored is being customized. Must not be null. 
     */
    public CustomOrderScoreQuery(Query subQuery) {
        this(subQuery, new ValueSourceQuery[0]);
    }

    /**
     * Create a CustomScoreQuery over input subQuery and a {@link ValueSourceQuery}.
     * @param subQuery the sub query whose score is being customized. Must not be null.
     * @param valSrcQuery a value source query whose scores are used in the custom score
     * computation. For most simple/convenient use case this would be a 
     * {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQuery}.
     * This parameter is optional - it can be null.
     */
    public CustomOrderScoreQuery(Query subQuery, ValueSourceQuery valSrcQuery) {
        this(subQuery, valSrcQuery != null ? // don't want an array that contains a single null.. 
                new ValueSourceQuery[]{valSrcQuery} : new ValueSourceQuery[0]);
    }

    /**
     * Create a CustomScoreQuery over input subQuery and a {@link ValueSourceQuery}.
     * @param subQuery the sub query whose score is being customized. Must not be null.
     * @param valSrcQueries value source queries whose scores are used in the custom score
     * computation. For most simple/convenient use case these would be 
     * {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQueries}.
     * This parameter is optional - it can be null or even an empty array.
     */
    public CustomOrderScoreQuery(Query subQuery, ValueSourceQuery... valSrcQueries) {
        this.subQuery = subQuery;
        this.valSrcQueries = valSrcQueries != null
                ? valSrcQueries : new ValueSourceQuery[0];
        if (subQuery == null) {
            throw new IllegalArgumentException("<subquery> must not be null!");
        }
    }

    /*(non-Javadoc) @see org.apache.lucene.search.Query#rewrite(org.apache.lucene.index.IndexReader) */
    @Override
    public Query rewrite(IndexReader reader) throws IOException {
        CustomOrderScoreQuery clone = null;
        final Query sq = subQuery.rewrite(reader);
        if (sq != subQuery) {
            clone = (CustomOrderScoreQuery) clone();
            clone.subQuery = sq;
        }
        for (int i = 0; i < valSrcQueries.length; i++) {
            final ValueSourceQuery v = (ValueSourceQuery) valSrcQueries[i].rewrite(reader);
            if (v != valSrcQueries[i]) {
                if (clone == null) {
                    clone = (CustomOrderScoreQuery) clone();
                }
                clone.valSrcQueries[i] = v;
            }
        }
        return (clone == null) ? this : clone;
    }

    /*(non-Javadoc) @see org.apache.lucene.search.Query#extractTerms(java.util.Set) */
    @Override
    public void extractTerms(Set<Term> terms) {
        subQuery.extractTerms(terms);
        for (int i = 0; i < valSrcQueries.length; i++) {
            valSrcQueries[i].extractTerms(terms);
        }
    }

    /*(non-Javadoc) @see org.apache.lucene.search.Query#clone() */
    @Override
    public Object clone() {
        CustomOrderScoreQuery clone = (CustomOrderScoreQuery) super.clone();
        clone.subQuery = (Query) subQuery.clone();
        clone.valSrcQueries = new ValueSourceQuery[valSrcQueries.length];
        for (int i = 0; i < valSrcQueries.length; i++) {
            clone.valSrcQueries[i] = (ValueSourceQuery) valSrcQueries[i].clone();
        }
        return clone;
    }

    /* (non-Javadoc) @see org.apache.lucene.search.Query#toString(java.lang.String) */
    @Override
    public String toString(String field) {
        StringBuilder sb = new StringBuilder(name()).append("(");
        sb.append(subQuery.toString(field));
        for (int i = 0; i < valSrcQueries.length; i++) {
            sb.append(", ").append(valSrcQueries[i].toString(field));
        }
        sb.append(")");
        sb.append(strict ? " STRICT" : "");
        return sb.toString() + ToStringUtils.boost(getBoost());
    }

    /** Returns true if <code>o</code> is equal to this. */
    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (!super.equals(o)) {
            return false;
        }
        if (getClass() != o.getClass()) {
            return false;
        }
        CustomOrderScoreQuery other = (CustomOrderScoreQuery) o;
        if (this.getBoost() != other.getBoost() || !this.subQuery.equals(other.subQuery) || this.strict != other.strict || this.valSrcQueries.length != other.valSrcQueries.length) {
            return false;
        }
        return Arrays.equals(valSrcQueries, other.valSrcQueries);
    }

    /** Returns a hash code value for this object. */
    @Override
    public int hashCode() {
        return (getClass().hashCode() + subQuery.hashCode() + Arrays.hashCode(valSrcQueries)) ^ Float.floatToIntBits(getBoost()) ^ (strict ? 1234 : 4321);
    }

    /**
     * Returns a {@link CustomScoreProvider} that calculates the custom scores
     * for the given {@link IndexReader}. The default implementation returns a default
     * implementation as specified in the docs of {@link CustomScoreProvider}.
     * @since 2.9.2
     */
    protected CustomScoreProvider getCustomScoreProvider(IndexReader reader) throws IOException {
        return new CustomScoreProvider(reader) {

            @Override
            public float customScore(int doc, float subQueryScore, float valSrcScore) {
                float result = subQueryScore * valSrcScore;
                return result;
            }
        };
    }

    //=========================== W E I G H T ============================
    private class CustomWeight extends Weight {

        Similarity similarity;
        Weight subQueryWeight;
        Weight[] valSrcWeights;
        boolean qStrict;

        public CustomWeight(Searcher searcher) throws IOException {
            this.similarity = getSimilarity(searcher);
            this.subQueryWeight = subQuery.createWeight(searcher);
            this.valSrcWeights = new Weight[valSrcQueries.length];
            for (int i = 0; i < valSrcQueries.length; i++) {
                this.valSrcWeights[i] = valSrcQueries[i].createWeight(searcher);
            }
            this.qStrict = strict;
        }

        /*(non-Javadoc) @see org.apache.lucene.search.Weight#getQuery() */
        @Override
        public Query getQuery() {
            return CustomOrderScoreQuery.this;
        }

        /*(non-Javadoc) @see org.apache.lucene.search.Weight#getValue() */
        @Override
        public float getValue() {
            return getBoost();
        }

        /*(non-Javadoc) @see org.apache.lucene.search.Weight#sumOfSquaredWeights() */
        @Override
        public float sumOfSquaredWeights() throws IOException {
            float sum = subQueryWeight.sumOfSquaredWeights();
            for (int i = 0; i < valSrcWeights.length; i++) {
                if (qStrict) {
                    valSrcWeights[i].sumOfSquaredWeights(); // do not include ValueSource part in the query normalization
                } else {
                    sum += valSrcWeights[i].sumOfSquaredWeights();
                }
            }
            sum *= getBoost() * getBoost(); // boost each sub-weight
            return sum;
        }

        /*(non-Javadoc) @see org.apache.lucene.search.Weight#normalize(float) */
        @Override
        public void normalize(float norm) {
            norm *= getBoost(); // incorporate boost
            subQueryWeight.normalize(norm);
            for (int i = 0; i < valSrcWeights.length; i++) {
                if (qStrict) {
                    valSrcWeights[i].normalize(1); // do not normalize the ValueSource part
                } else {
                    valSrcWeights[i].normalize(norm);
                }
            }
        }

        @Override
        public Scorer scorer(IndexReader reader, boolean scoreDocsInOrder, boolean topScorer) throws IOException {
            // Pass true for "scoresDocsInOrder", because we
            // require in-order scoring, even if caller does not,
            // since we call advance on the valSrcScorers.  Pass
            // false for "topScorer" because we will not invoke
            // score(Collector) on these scorers:
            Scorer subQueryScorer = subQueryWeight.scorer(reader, true, false);
            if (subQueryScorer == null) {
                return null;
            }
            Scorer[] valSrcScorers = new Scorer[valSrcWeights.length];
            for (int i = 0; i < valSrcScorers.length; i++) {
                valSrcScorers[i] = valSrcWeights[i].scorer(reader, true, topScorer);
            }
            return new CustomScorer(similarity, reader, this, subQueryScorer, valSrcScorers);
        }

        @Override
        public Explanation explain(IndexReader reader, int doc) throws IOException {
            Explanation explain = doExplain(reader, doc);
            return explain == null ? new Explanation(0.0f, "no matching docs") : explain;
        }

        private Explanation doExplain(IndexReader reader, int doc) throws IOException {
            Explanation subQueryExpl = subQueryWeight.explain(reader, doc);
            if (!subQueryExpl.isMatch()) {
                return subQueryExpl;
            }
            // match
            Explanation[] valSrcExpls = new Explanation[valSrcWeights.length];
            for (int i = 0; i < valSrcWeights.length; i++) {
                valSrcExpls[i] = valSrcWeights[i].explain(reader, doc);
            }
            Explanation customExp = CustomOrderScoreQuery.this.getCustomScoreProvider(reader).customExplain(doc, subQueryExpl, valSrcExpls);
            float sc = getValue() * customExp.getValue();
            Explanation res = new ComplexExplanation(
                    true, sc, CustomOrderScoreQuery.this.toString() + ", product of:");
            res.addDetail(customExp);
            res.addDetail(new Explanation(getValue(), "queryBoost")); // actually using the q boost as q weight (== weight value)
            return res;
        }

        @Override
        public boolean scoresDocsOutOfOrder() {
            return false;
        }
    }

    //=========================== S C O R E R ============================
    /**
     * A scorer that applies a (callback) function on scores of the subQuery.
     */
    private class CustomScorer extends Scorer {

        private final float qWeight;
        private Scorer subQueryScorer;
        private Scorer[] valSrcScorers;
        private final CustomScoreProvider provider;
        private float vScores[]; // reused in score() to avoid allocating this array for each doc 

        // constructor
        private CustomScorer(Similarity similarity, IndexReader reader, CustomWeight w,
                Scorer subQueryScorer, Scorer[] valSrcScorers) throws IOException {
            super(similarity, w);
            this.qWeight = w.getValue();
            this.subQueryScorer = subQueryScorer;
            this.valSrcScorers = valSrcScorers;
            this.vScores = new float[valSrcScorers.length];
            this.provider = CustomOrderScoreQuery.this.getCustomScoreProvider(reader);
        }

        @Override
        public int nextDoc() throws IOException {
            int doc = subQueryScorer.nextDoc();
            if (doc != NO_MORE_DOCS) {
                for (int i = 0; i < valSrcScorers.length; i++) {
                    valSrcScorers[i].advance(doc);
                }
            }
            return doc;
        }

        @Override
        public int docID() {
            return subQueryScorer.docID();
        }

        /*(non-Javadoc) @see org.apache.lucene.search.Scorer#score() */
        @Override
        public float score() throws IOException {
            for (int i = 0; i < valSrcScorers.length; i++) {
                vScores[i] = valSrcScorers[i].score();
            }
            return qWeight * provider.customScore(subQueryScorer.docID(), subQueryScorer.score(), vScores);
        }

        @Override
        public int advance(int target) throws IOException {
            int doc = subQueryScorer.advance(target);
            if (doc != NO_MORE_DOCS) {
                for (int i = 0; i < valSrcScorers.length; i++) {
                    valSrcScorers[i].advance(doc);
                }
            }
            return doc;
        }
    }

    @Override
    public Weight createWeight(Searcher searcher) throws IOException {
        return new CustomWeight(searcher);
    }

    /**
     * Checks if this is strict custom scoring.
     * In strict custom scoring, the ValueSource part does not participate in weight normalization.
     * This may be useful when one wants full control over how scores are modified, and does 
     * not care about normalizing by the ValueSource part.
     * One particular case where this is useful if for testing this query.   
     * <P>
     * Note: only has effect when the ValueSource part is not null.
     */
    public boolean isStrict() {
        return strict;
    }

    /**
     * Set the strict mode of this query. 
     * @param strict The strict mode to set.
     * @see #isStrict()
     */
    public void setStrict(boolean strict) {
        this.strict = strict;
    }

    /**
     * A short name of this query, used in {@link #toString(String)}.
     */
    public String name() {
        return "custom";
    }
}
